import os
import json
from collections import defaultdict

class UnionFind:
    def __init__(self):
        self.parent = {}
    
    def find(self, u):
        if u not in self.parent:
            self.parent[u] = u
        while self.parent[u] != u:
            self.parent[u] = self.parent[self.parent[u]]  # 路径压缩
            u = self.parent[u]
        return u
    
    def union(self, u, v):
        root_u = self.find(u)
        root_v = self.find(v)
        if root_u != root_v:
            self.parent[root_v] = root_u

def parse_obj(obj_path):
    vertices = []
    tex_coords = []
    normals = []
    faces = []
    other_lines = []

    with open(obj_path, 'r') as f:
        for line in f:
            stripped = line.strip()
            if not stripped:
                continue
            if stripped.startswith('v '):
                parts = stripped.split()
                vertices.append(list(map(float, parts[1:4])))
            elif stripped.startswith('vt'):
                parts = stripped.split()
                tex_coords.append(list(map(float, parts[1:])))
            elif stripped.startswith('vn'):
                parts = stripped.split()
                normals.append(list(map(float, parts[1:])))
            elif stripped.startswith('f'):
                parts = stripped.split()
                face = []
                for part in parts[1:]:
                    indices = part.split('/')
                    v_idx = int(indices[0]) - 1
                    vt_idx = int(indices[1]) - 1 if len(indices) > 1 and indices[1] else None
                    vn_idx = int(indices[2]) - 1 if len(indices) > 2 and indices[2] else None
                    face.append((v_idx, vt_idx, vn_idx))
                faces.append(face)
            else:
                other_lines.append(line.rstrip())
    return vertices, tex_coords, normals, faces, other_lines

def split_obj(obj_path, output_prefix="output"):
    vertices, tex_coords, normals, faces, other_lines = parse_obj(obj_path)
    
    uf = UnionFind()
    for face in faces:
        if face:
            first_v = face[0][0]
            for vertex in face[1:]:
                uf.union(first_v, vertex[0])
    
    groups = defaultdict(list)
    for v_idx in range(len(vertices)):
        root = uf.find(v_idx)
        groups[root].append(v_idx)
    
    metadata = {
        "num_original_vertices": len(vertices),
        "num_original_tex_coords": len(tex_coords),
        "num_original_normals": len(normals),
        "groups": []
    }
    
    for group_idx, (root, v_indices) in enumerate(groups.items(), 1):
        group_v = set(v_indices)
        group_faces = [face for face in faces if face and face[0][0] in group_v]
        
        vt_indices = set()
        vn_indices = set()
        for face in group_faces:
            for vertex in face:
                if vertex[1] is not None:
                    vt_indices.add(vertex[1])
                if vertex[2] is not None:
                    vn_indices.add(vertex[2])
        
        # 计算原始顶点坐标的中心点
        original_vertices = [vertices[idx] for idx in v_indices]
        num_vertices = len(original_vertices)
        if num_vertices == 0:
            cx, cy, cz = 0.0, 0.0, 0.0
        else:
            cx = sum(v[0] for v in original_vertices) / num_vertices
            cy = sum(v[1] for v in original_vertices) / num_vertices
            cz = sum(v[2] for v in original_vertices) / num_vertices
        
        # 创建平移后的顶点、纹理坐标和法线
        new_vertices = [[v[0] - cx, v[1] - cy, v[2] - cz] for v in original_vertices]
        new_tex_coords = [tex_coords[idx] for idx in vt_indices] if tex_coords else []
        new_normals = [normals[idx] for idx in vn_indices] if normals else []
        
        # 构建输出内容
        output = []
        output.extend(other_lines)
        
        for v in new_vertices:
            output.append(f"v {' '.join(map(str, v))}")
        for vt in new_tex_coords:
            output.append(f"vt {' '.join(map(str, vt))}")
        for vn in new_normals:
            output.append(f"vn {' '.join(map(str, vn))}")
        
        # 创建索引映射
        v_index_map = {orig_idx: i+1 for i, orig_idx in enumerate(v_indices)}
        vt_index_map = {orig_idx: i+1 for i, orig_idx in enumerate(sorted(vt_indices))} if vt_indices else {}
        vn_index_map = {orig_idx: i+1 for i, orig_idx in enumerate(sorted(vn_indices))} if vn_indices else {}
        
        for face in group_faces:
            face_line = []
            for vertex in face:
                v_orig = vertex[0]
                vt_orig = vertex[1]
                vn_orig = vertex[2]
                
                v_new = v_index_map[v_orig]
                vt_new = vt_index_map.get(vt_orig, '') if vt_orig is not None else ''
                vn_new = vn_index_map.get(vn_orig, '') if vn_orig is not None else ''
                
                parts = [str(v_new)]
                if vt_new or vn_new:
                    parts.append(str(vt_new) if vt_new else '')
                    if vn_new:
                        parts.append(str(vn_new))
                face_part = '/'.join(parts)
                face_line.append(face_part)
            output.append(f"f {' '.join(face_line)}")
        
        output_path = f"{output_prefix}_{group_idx}.obj"
        with open(output_path, 'w') as f:
            f.write('\n'.join(output))
        print(f"Created: {output_path}")
        
        # 保存元数据信息
        group_metadata = {
            "output_file": output_path,
            "transform": {
                "translate": [cx, cy, cz]
            },
            "v_indices": v_indices,
            "vt_indices": sorted(vt_indices) if vt_indices else [],
            "vn_indices": sorted(vn_indices) if vn_indices else []
        }
        metadata["groups"].append(group_metadata)
    
    # 保存元数据到JSON文件
    metadata_path = f"{output_prefix}_metadata.json"
    with open(metadata_path, 'w') as f:
        json.dump(metadata, f, indent=4)
    print(f"Metadata saved to {metadata_path}")

if __name__ == "__main__":
    input_obj = "qiaoduntai.obj"
    split_obj(input_obj)